import ilqr_controller
import math
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.patches import Rectangle, Polygon

def cal_vehicle_pts(x,y,theta,width,long):
    
    pts=[[-long/2,width/2],[long/2,width/2],[long/2,-width/2],[-long/2,-width/2]]
    pt_list_res=[]
    rotation_matrix = np.array([
    [np.cos(theta), -np.sin(theta)],
    [np.sin(theta), np.cos(theta)]])
    for pt in pts:
        local_points = np.array([pt[0], pt[1]])
        rotated_point = np.dot(rotation_matrix, local_points)
        world_point = rotated_point + np.array([x, y])
        pt_list_res.append((world_point[0],world_point[1]))
    return pt_list_res


def Draw(result_proto,traj):
    fig, ax = plt.subplots()
    ax.set_aspect('equal')
    ax.set_xlim(-30, 30)
    ax.set_ylim(-30, 30)

    control_traj=result_proto.ilqr_iteration_datas[len(result_proto.ilqr_iteration_datas)-1]
    for i,traj_pt in enumerate(control_traj.states):
        opacity = (i + 1) / len(control_traj.states)  # 让透明度随着时间增加而增加
        width = 1
        long = 2
        print("traj_theta=",traj_pt.theta)
        vehicle_pts=cal_vehicle_pts(traj_pt.x , traj_pt.y,traj_pt.theta,1,2)
        rect = Polygon(vehicle_pts,
                        closed=True, color='blue', alpha=1-opacity)  
        ax.add_patch(rect)
        rect.set_visible(True)

    traj_plot = ax.scatter([], [], s=2)
    traj_plot_data=np.zeros((len(traj),2),float)
    for k,traj_pt in enumerate(traj):
        traj_plot_data[k]=[traj_pt[0],traj_pt[1]]
    traj_plot.set_offsets(traj_plot_data)
    traj_plot.set_color('black')
    traj_plot.set_visible(True)

    warm_start_plot = ax.scatter([], [], s=2)
    warm_start_data=np.zeros((len(result_proto.warm_start_data.states),2),float)
    for k,warm_start_pt in enumerate(result_proto.warm_start_data.states):
        warm_start_data[k]=[warm_start_pt.x,warm_start_pt.y]
    warm_start_plot.set_offsets(warm_start_data)
    warm_start_plot.set_color('red')
    warm_start_plot.set_visible(True)



    def init():
        return []


    #最后一帧用来清除
    def update(frame):


        return 

    ani = FuncAnimation(fig, update, frames=1, init_func=init, blit=False,interval=1000)

    plt.show()











ilqr=ilqr_controller.IlqrController()

try:
    preFile=np.loadtxt('trajeo.txt')
except:
    print('traj dont exist')


rows,cols=preFile.shape
x = preFile[:, 0]
y = preFile[:, 1]
h = preFile[:, 2]
v = preFile[:, 3]

t=np.zeros((x.shape[0]))
if cols<5:
    for i in range(x.shape[0]):
        if i==0:
            t[i]=0.0
        else:
            t[i]=t[i-1]+(math.sqrt((x[i]-x[i-1])**2+(y[i]-y[i-1])**2))/((v[i]+v[i-1])/2)
else:
    t=preFile[:,4]

ilqr.update_trajectory(x, y, h, v, t)

import time
t1=time.time()
ilqr.update_ego_state(335900.647 ,758497.729 ,356.967 ,0.327,0.0)
ilqr.process()
t2=time.time()
print("consume time = ",t2-t1)
result_proto=ilqr.result_proto
Draw(result_proto,ilqr.traj_base_ego)
